import SwiftCompilerPlugin
import SwiftSyntax
import SwiftSyntaxBuilder
import SwiftSyntaxMacros

/// Implementation of the `stringify` macro, which takes an expression
/// of any type and produces a tuple containing the value of that expression
/// and the source code that produced the value. For example
///
///     #stringify(x + y)
///
///  will expand to
///
///     (x + y, "x + y")
public struct StringifyMacro: ExpressionMacro {
    public static func expansion(
        of node: some FreestandingMacroExpansionSyntax,
        in context: some MacroExpansionContext
    ) -> ExprSyntax {
        guard let argument = node.arguments.first?.expression else {
            fatalError("compiler bug: the macro does not have any arguments")
        }

        return "(\(argument), \(literal: argument.description))"
    }
}

public struct DefineKeysMacro: DeclarationMacro {
    
    public static func expansion(of node: some FreestandingMacroExpansionSyntax, in context: some MacroExpansionContext) throws -> [DeclSyntax] {
        let idList: [String] = try node.arguments.map { les in
            guard let stringLiteral = les.expression.as(StringLiteralExprSyntax.self), let key = stringLiteral.segments.first?.description.trimmingCharacters(in: .init(charactersIn: "\"")) else {
                throw MacroExpansionErrorMessage("#defineKeys 参数必须为字符串字面量")
            }
            
            return key
        }
        
        return idList.map { key in
            DeclSyntax("let \(raw: key) = \"\(raw: key)\"")
        }
    }
    
}

public struct AddIdMacro: MemberMacro {
    
    public static func expansion(of node: AttributeSyntax, providingMembersOf declaration: some DeclGroupSyntax, in context: some MacroExpansionContext) throws -> [DeclSyntax] {
        return [DeclSyntax(stringLiteral: "var id = UUID().uuidString")]
    }
    
}

public struct CodableSubclassMacro: MemberMacro {
    
    public static func expansion(of node: AttributeSyntax, providingMembersOf declaration: some DeclGroupSyntax, in context: some MacroExpansionContext) throws -> [DeclSyntax] {
        // 1、验证是否是类声明
        guard let classDecl = declaration.as(ClassDeclSyntax.self) else {
            throw MacroExpansionErrorMessage("@CodableSubclass 只能用于类")
        }
        // 2、验证是否有父类
        guard let inheritanceClause = declaration.inheritanceClause, inheritanceClause.inheritedTypes.contains(where: { type in
            type.type.trimmedDescription == "BaseModel" || type.type.trimmedDescription.contains("Codable")
        }) else {
            throw MacroExpansionErrorMessage("@CodableSubclass 需要继承自 Codable 父类")
        }
        
        // 3、收集所有存储属性
        let storedProperties = classDecl.memberBlock.members.compactMap({ $0.decl.as(VariableDeclSyntax.self) }).filter({ $0.bindingSpecifier.text == "var" }).flatMap({ $0.bindings }).compactMap { binding -> String? in
            guard let pattern = binding.pattern.as(IdentifierPatternSyntax.self) else {
                return nil
            }
            
            return pattern.identifier.text
        }
        
        // 4、生成 CodingKeys枚举
        let codingKeysEnum = try EnumDeclSyntax("private enum CodingKeys: String, CodingKey", membersBuilder: {
            for property in storedProperties {
                "case \(raw: property)"
            }
        })
        
        // 5、生成 init(from:) 方法
        let initializer = try InitializerDeclSyntax("required init(from decoder: Decoder) throws", bodyBuilder: {
            // 调用父类解码器
            "try super.init(from: decoder)"
            
            // 创建容器
            "let container = try decoder.container(keyedBy: CodingKeys.self)"
            
            // 解码每个属性
            for property in storedProperties {
                "self.\(raw: property) = try container.decode(\(raw: getTypeName(for: property, in: declaration)).self, forKey: .\(raw: property))"
            }
        })
        
        return [DeclSyntax(codingKeysEnum), DeclSyntax(initializer)]
    }
    
    private static func getTypeName(for property: String, in declaration: some DeclGroupSyntax) -> String {
        for member in declaration.memberBlock.members {
            guard let varDecl = member.decl.as(VariableDeclSyntax.self) else {
                continue
            }
            
            for binding in varDecl.bindings {
                guard let identifierPattern = binding.pattern.as(IdentifierPatternSyntax.self), identifierPattern.identifier.text == property else {
                    continue
                }
                
                if let typeAnnotation = binding.typeAnnotation {
                    return typeAnnotation.type.trimmedDescription
                }
            }
        }
        
        return "Any"
    }
    
}

public struct TestBodyMacro: BodyMacro {
    public static func expansion(of node: SwiftSyntax.AttributeSyntax, providingBodyFor declaration: some SwiftSyntax.DeclSyntaxProtocol & SwiftSyntax.WithOptionalCodeBlockSyntax, in context: some SwiftSyntaxMacros.MacroExpansionContext) throws -> [SwiftSyntax.CodeBlockItemSyntax] {
        let bodyText = "print(\"你好\")"
        let exitItem = CodeBlockItemSyntax(stringLiteral: bodyText)
        return [exitItem]
    }
    
}

public struct DescriptionMacro: MemberMacro {
    
    public static func expansion(of node: AttributeSyntax, providingMembersOf declaration: some DeclGroupSyntax, in context: some MacroExpansionContext) throws -> [DeclSyntax] {
        guard let classDecl = declaration.as(ClassDeclSyntax.self) else {
            fatalError("compiler bug: unknown error")
        }
        
        guard let inheritedTypes = declaration.inheritanceClause?.inheritedTypes, inheritedTypes.compactMap({ $0.type.as(IdentifierTypeSyntax.self) }).contains(where: { $0.name.text == "CustomStringConvertible" }) else {
            fatalError("compiler bug: need `CustomStringConvertible` conformance")
        }
        
        let className = classDecl.name.text
        let variables = classDecl.memberBlock.members.compactMap({ $0.decl.as(VariableDeclSyntax.self)?.bindings }).compactMap({ $0.first?.pattern.as(IdentifierPatternSyntax.self)?.identifier.text }).map({ "\($0): \\(\($0))" }).joined(separator: ", ")
        
        let ds = DeclSyntax("""
        var description: String {
            "\(raw: className)(\(raw: variables))"
        }
        """)
        return [ds]
    }
    
}

@main
struct MyFirstMacroPlugin: CompilerPlugin {
    let providingMacros: [Macro.Type] = [
        StringifyMacro.self,
        DefineKeysMacro.self,
        AddIdMacro.self,
        CodableSubclassMacro.self,
        TestBodyMacro.self,
        DescriptionMacro.self
    ]
}
